package com.example.datasourcedemo.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class DynamicDataSource extends AbstractRoutingDataSource {
    private static final Logger log = LoggerFactory.getLogger(DynamicDataSource.class);
    private static final ThreadLocal<String> dataSourceKey = new InheritableThreadLocal<>();

    public static void setDataSourceKey(String key) {
        log.info("设置数据源 Key: {}", key);
        dataSourceKey.set(key);
    }

    public static void clearDataSourceKey() {
//        log.info("清除数据源 Key: {}", dataSourceKey.get());
//        dataSourceKey.remove();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String key = dataSourceKey.get();
        log.debug("当前数据源 Key: {}", key);
        return key;
    }

    private final Map<Object, Object> dynamicDataSources = new ConcurrentHashMap<>();

    public void addDataSource(String key, DataSource dataSource) {
        dynamicDataSources.put(key, dataSource);
        super.setTargetDataSources(dynamicDataSources);
        super.afterPropertiesSet();
        log.info("已注册数据源 Key: {}", key);
    }

    public boolean containsDataSource(String key) {
        return dynamicDataSources.containsKey(key);
    }
}
